"""基于梯度卷积的计算方法"""

import numpy as np
import cv2
from typing import Iterable
from utils.reader import Projection
from scipy.interpolate import RectBivariateSpline
from scipy.ndimage import gaussian_filter
import sklearn.cluster as cluster
import tqdm


def dem_interpolation(dem, cntx, cnty):
    """
    输入3 x 3的DEM矩阵和当前中心坐标，返回中心处的插值结果，用双线性插值
    """
    X, Y = np.arange(5), np.arange(5)
    dem_ = np.pad(dem, ((2, 2), (2, 2)), "edge")
    if isinstance(cntx, np.ndarray):
        z = []
        for x, y in zip(cntx, cnty):
            Z = dem_[int(x) : int(x) + 5, int(y) : int(y) + 5]
            f = RectBivariateSpline(X, Y, Z)
            z.append(f(2, 2)[0, 0])
        return np.array(z)
    else:
        Z = dem_[int(cntx) : int(cntx) + 5, int(cnty) : int(cnty) + 5]
        f = RectBivariateSpline(X, Y, Z)
        return f(2, 2)[0, 0]


def dem_crater_filter(
    nac_prj_: Projection,
    dem_prj_: Projection,
    lat_lon_r_iter: Iterable[tuple[float, float, float]],
    depth_th=8,
    fluctuation_th=0.2,
):
    dem = dem_prj_.image
    grad_x = cv2.Sobel(dem, cv2.CV_64F, 1, 0, ksize=3)
    grad_y = cv2.Sobel(dem, cv2.CV_64F, 0, 1, ksize=3)
    mag = cv2.magnitude(grad_x, grad_y)
    for lat, lon, r in lat_lon_r_iter:
        R = 1.5 * r
        nac_x_cnt, nac_y_cnt = nac_prj_.forward(lat, lon)
        left, top = nac_x_cnt - R, nac_y_cnt - R
        right, down = nac_x_cnt + R, nac_y_cnt + R
        # crater_nac = nac_prj_.image[int(left) : int(right), int(top) : int(down)]
        # x_cnt, y_cnt = dem_prj_.forward(lat, lon)
        # left, top = x_cnt - r, y_cnt - r - 12
        # right, down = x_cnt + r, y_cnt + r - 12
        left, top = dem_prj_.forward(*nac_prj_.inverse(left, top))
        right, down = dem_prj_.forward(*nac_prj_.inverse(right, down))
        left = max(0, left)
        top = max(0, top)
        right = min(dem.shape[0], right)
        down = min(dem.shape[1], down)
        # crater_dem = dem[int(left) : int(right), int(top) : int(down)]
        # dem_ = (
        #     (crater_dem - crater_dem.min())
        #     / (crater_dem.max() - crater_dem.min())
        #     * 255
        # )
        # nac = (
        #     (crater_nac - crater_nac.min())
        #     / (crater_nac.max() - crater_nac.min())
        #     * 255
        # )
        # cv2.imwrite("nac.png", nac)
        # cv2.imwrite("dem.png", dem_)
        dem_patch = dem[int(left) : int(right), int(top) : int(down)]
        mag_patch = mag[int(left) : int(right), int(top) : int(down)]
        # 选出梯度较小且高度较小的点，作为圆心点
        x_y = np.argwhere(
            (dem_patch < np.percentile(dem_patch, 10))
            & (mag_patch < np.percentile(mag_patch, 20))
        )
        if len(x_y) == 0 or x_y.std(axis=0).sum() > 4:
            continue
        # 用聚类算法把可能的不同中心都分出去
        # db = cluster.DBSCAN(eps=5, min_samples=1).fit(x_y)
        # for label in np.unique(db.labels_):
        # 针对每一个中心，将当前检测框的中心平移到这里
        dem_x_cnt, dem_y_cnt = x_y.mean(axis=0)
        # r = sum(dem_patch.shape) / 4
        dem_x_cnt, dem_y_cnt = dem_x_cnt + left, dem_y_cnt + top
        # cv2.circle(nac, (int(y - top), int(x - left)), 1, (255, 255, 255), 1)
        # cv2.circle(dem_, (int(y - top), int(x - left)), 1, (255, 255, 255), 1)
        # if x < 2 or x >= dem.shape[0] - 2:
        #     continue
        # if y < 2 or y >= dem.shape[1] - 2:
        #     continue
        # z = dem_interpolation(dem, x, y)
        rim_z = circle_extract(dem, dem_x_cnt, dem_y_cnt, r)
        if len(rim_z) < 6:
            continue
        # for x_, y_ in rim_z[:, :2]:
        #     cv2.circle(nac, (int(y_ - top), int(x_ - left)), 1, (0, 0, 255), 1)
        #     cv2.circle(dem_, (int(y_ - top), int(x_ - left)), 1, (0, 0, 255), 1)
        # cntx, cnty, r = radii_fit(rim_z[:, 0], rim_z[:, 1])
        # cv2.circle(dem_, (int(cnty - top), int(cntx - left)), int(r), (0, 0, 255), 1)
        # cv2.circle(nac, (int(cnty - top), int(cntx - left)), int(r), (0, 0, 255), 1)
        # cv2.imwrite("nac.png", nac)
        # cv2.imwrite("dem.png", dem_)
        # if r is None or np.abs(2 * r / sum(dem.shape)) > 2:
        #     # 半径变化过于剧烈，不予考虑
        #     continue
        lat, lon = dem_prj_.inverse(dem_x_cnt, dem_y_cnt)
        # nac_ = (
        #     (nac_prj_.image - nac_prj_.image.min())
        #     / (nac_prj_.image.max() - nac_prj_.image.min())
        #     * 255
        # )
        # cv2.circle(nac_, (int(y ), int(x)), round(r), (0, 0, 255), 1)
        # cv2.imwrite("nac_.png", nac_)
        nac_x, nac_y = nac_prj_.forward(*dem_prj_.inverse(rim_z[:, 0], rim_z[:, 1]))
        nac_x_cnt_esti, nac_y_cnt_esti, r_ = radii_fit(nac_x, nac_y)
        if r_ is None:
            continue
        lat_, lon_ = nac_prj_.inverse(nac_x_cnt_esti, nac_y_cnt_esti)
        # 排除错误的拟合结果
        if (
            np.sqrt(
                (nac_x_cnt_esti - nac_x_cnt) ** 2 + (nac_y_cnt_esti - nac_y_cnt) ** 2
            )
            > r
            or np.abs(r - r_) > 0.5 * r
        ):
            continue
        dem_x, dem_y = dem_prj_.forward(lat_, lon_)
        z = dem[int(dem_x), int(dem_y)]
        if (rim_z[:, 2].mean() - z) * 100 / r_ > depth_th and rim_z[
            :, 2
        ].std() < fluctuation_th * r_:
            yield lat_, lon_, r_


def circle_extract(dem, cntx, cnty, r):
    rim_z = []
    for theta in np.arange(0, 2 * np.pi, np.pi / 10):
        ri = np.arange(0.2 * r, 1.2 * r, 1)
        u = cntx + ri * np.cos(theta)
        v = cnty + ri * np.sin(theta)
        u, v = np.round(u).astype(int), np.round(v).astype(int)
        ind = (u >= 0) & (u < dem.shape[0] - 1) & (v >= 0) & (v < dem.shape[1] - 1)
        u, v = u[ind], v[ind]
        # r_z = dem_interpolation(dem, u[ind], v[ind])
        r_z = dem[u, v]
        # 选取最佳的点作为边缘点，即梯度从正变为负或者变为零的点
        if len(r_z) < 4:
            continue
        # r_z = gaussian_filter(r_z, 2)
        grad = np.diff(np.pad(r_z, ((1, 0)), "edge"))
        cross = np.where((grad[:-1] > 0) & (grad[1:] <= 0))[0]
        if len(cross) == 0:
            continue
        # 保留最高点
        cross = cross[np.argmax(r_z[cross])]
        # cross = cross[0]
        rim_z.append(np.array((u[cross], v[cross], r_z[cross])))
    return np.array(rim_z)


def radii_fit(x_sample: np.ndarray, y_sample: np.ndarray):
    """
    用八点法拟合圆形，即在DEM上获取八个方向上的最高点（或近似最高点），转译为NAC上的坐标，完成圆形拟合。
    """
    X = np.vstack([x_sample, y_sample, np.ones_like(x_sample)]).T
    Y = np.vstack([x_sample * x_sample + y_sample * y_sample]).T
    a, b, c = np.linalg.lstsq(X, Y, rcond=None)[0]
    cntx = a / 2
    cnty = b / 2
    r = np.sqrt(c + cntx**2 + cnty**2)
    return cntx.item(), cnty.item(), r.item()


def ransac_circle_fit(x, y, t=4, k=100, th=20):
    sample = np.random.choice(range(len(x)), t, replace=False)
    best_num = 0
    best_choice = None
    for _ in range(k):
        cntx, cnty, r = radii_fit(x[sample], y[sample])
        err = np.abs((x - cntx) ** 2 + (y - cnty) ** 2 - r * r)
        inliers = err < th
        if inliers.sum() > best_num:
            best_num = inliers.sum()
            best_choice = inliers
            if best_num > t:
                sample = np.arange(len(x))[inliers]
            else:
                sample = np.random.choice(range(len(x)), t, replace=False)
        else:
            sample = np.random.choice(range(len(x)), t, replace=False)
    if best_choice is None:
        return None, None, None
    else:
        cntx, cnty, r = radii_fit(x[best_choice], y[best_choice])
        cntx_, cnty_, r_ = radii_fit(x, y)
        if np.sqrt((cntx - cntx_) ** 2 + (cnty - cnty_) ** 2 + (r - r_) ** 2) > 100:
            return None, None, None
        return cntx, cnty, r


def crater_filter(
    craters: np.ndarray,
    dem_prj: Projection,
    nac_prj: Projection,
    depth_th=0.5,
    fluctuation_th=0.8,
):
    craters = np.array(craters)
    # 仅使用中心作为聚类的特征
    labels = cluster.DBSCAN(eps=10, min_samples=2).fit(craters).labels_
    # 对于每个类别，拟合一个圆形
    for cntx, cnty, r in tqdm.tqdm(craters[labels == -1], desc="Crater Filter"):
        # x, y = dem_prj.forward(*nac_prj.inverse(cntx, cnty))
        # if x < 0 or x >= dem_prj.image.shape[0] - 1:
        #     continue
        # if y < 0 or y >= dem_prj.image.shape[1] - 1:
        #     continue
        # # z = dem_interpolation(dem_prj.image, x, y)
        # z = dem_prj.image[int(x), int(y)]
        # rim_z = circle_extract(dem_prj.image, x, y, r)
        # if len(rim_z) < 6:
        #     continue
        # x, y = nac_prj.forward(*dem_prj.inverse(rim_z[:, 0], rim_z[:, 1]))
        # cntx, cnty, r = radii_fit(x, y)
        # if r is None:
        #     continue
        lat, lon = nac_prj.inverse(cntx, cnty)
        # rim_z = circle_extract(
        #     dem_prj.image, *dem_prj.forward(*nac_prj.inverse(cntx, cnty)), r
        # )
        # rim_z = rim_z[:, 2]
        # if rim_z.mean() - z > depth_th and rim_z.std() < fluctuation_th:
        yield lat, lon, r
    for label in tqdm.tqdm(np.unique(labels)[1:], desc="Crater Filter"):
        # 标签为-1的是DBSCAN算法输出的噪声点，这里应当认为是单个类别点，留在后面来处理
        x = craters[labels == label][:, 0].mean()
        y = craters[labels == label][:, 1].mean()
        d = np.sqrt(
            (craters[labels == label][:, 0] - x) ** 2
            + (craters[labels == label][:, 1] - y) ** 2
        )
        r = (craters[labels == label][:, 2] + d).max()
        x, y = dem_prj.forward(*nac_prj.inverse(x, y))
        rim_z = circle_extract(dem_prj.image, x, y, r)
        if len(rim_z) < 6:
            continue
        x, y = nac_prj.forward(*dem_prj.inverse(rim_z[:, 0], rim_z[:, 1]))
        cntx, cnty, r = ransac_circle_fit(x, y)
        if r is None:
            continue
        rim_z = circle_extract(
            dem_prj.image, *dem_prj.forward(*nac_prj.inverse(cntx, cnty)), r
        )
        if len(rim_z) < 6:
            continue
        rim_z = rim_z[:, 2]
        lat, lon = nac_prj.inverse(cntx, cnty)
        x, y = dem_prj.forward(lat, lon)
        if x < 0 or x >= dem_prj.image.shape[0] - 1:
            continue
        if y < 0 or y >= dem_prj.image.shape[1] - 1:
            continue
        # z = dem_interpolation(dem_prj.image, x, y)
        z = dem_prj.image[int(x), int(y)]
        if rim_z.mean() - z > depth_th and rim_z.std() < fluctuation_th:
            yield lat, lon, r


if __name__ == "__main__":
    import yaml
    from utils.reader import from_TIFF, from_IMG

    config_path = "/home/a804_cbf/Code/luner_crater/config/2080x4/HG-HV.yaml"
    with open(config_path, "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
        nac_config = config["nac"]
        dem_config = config["dem"]
        subfig_config = config["subfig"]
    dem_prj = from_TIFF(**dem_config, **subfig_config)
    nac_prj = from_IMG(nac_config, **subfig_config)
    nac = nac_prj.image.copy()
    dem = dem_prj.image.copy()
    nac = (nac - nac.min()) / (nac.max() - nac.min()) * 255
    dem[dem == dem.min()] = dem.max()
    dem = (dem - dem.min()) / (dem.max() - dem.min()) * 255
    nac = cv2.cvtColor(nac, cv2.COLOR_GRAY2BGR)
    dem = cv2.cvtColor(dem, cv2.COLOR_GRAY2BGR)
    with open("/disk527/sdb1/a804_cbf/datasets/chang_e/crater_catalog.txt", "r") as f:
        for line in f:
            lat, lon, r = map(float, line.split(","))
            x1, y1 = nac_prj.forward(lat, lon)
            x2, y2 = dem_prj.forward(lat, lon)
            cv2.circle(nac, (round(y1), round(x1)), round(r), (0, 0, 255), 1)
            cv2.ellipse(
                dem,
                (round(y2), round(x2)),
                (round(r), round(r * np.cos(lat / 180 * np.pi))),
                0,
                0,
                360,
                (0, 0, 255),
                1,
            )
    cv2.imwrite("nac.png", nac)
    cv2.imwrite("dem.png", dem)
